# Definition for a binary tree node.
from collections import defaultdict
from typing import List


class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None


class Solution:
    def __init__(self):
        self.father = {}  # 字典储存各个节点的父节点
        self.ans = []  # 储存结果

    def find_father(self, root: TreeNode):
        """
        递归搜索各个节点的父节点，储存在类的father字典中
        :param root:  当前节点
        :return:
        """
        if root.left is not None:
            self.father[root.left.val] = root
            self.find_father(root.left)
        if root.right is not None:
            self.father[root.right.val] = root
            self.find_father(root.right)

    def find_ans(self, root: TreeNode, former: TreeNode, distance: int, k: int):
        """
        搜索和目标节点距离为k的节点
        :param root: 当前节点
        :param former: 上一个节点
        :param distance:  当前和目标节点的距离
        :param k: 目标距离
        :return:
        """
        if root is None:
            return
        if distance == k:
            self.ans.append(root.val)
            return
        if root.left != former:
            self.find_ans(root.left, root, distance + 1, k)
        if root.right != former:
            self.find_ans(root.right, root, distance + 1, k)
        if self.father[root.val] != former:
            self.find_ans(self.father[root.val], root, distance + 1, k)

    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        # 从根节点出发，搜索、记录树上每个节点的父节点
        self.father[root.val] = None # 需要给根节点手动赋值
        self.find_father(root)
        self.find_ans(target, None, 0, k)
        return self.ans


class Solution2:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        ans = []
        mp = defaultdict(TreeNode)

        def getParents(node):
            if node.left:
                mp[node.left.val] = node
                getParents(node.left)
            if node.right:
                mp[node.right.val] = node
                getParents(node.right)

        def findNode(node, fnode, depth, k):
            if node is None:
                return
            if depth == k:
                ans.append(node.val)
                return
            if node.left != fnode:
                findNode(node.left, node, depth + 1, k)
            if node.right != fnode:
                findNode(node.right, node, depth + 1, k)
            if mp[node.val] != fnode:
                findNode(mp[node.val], node, depth + 1, k)

        mp[root.val] = None
        getParents(root)
        findNode(target, None, 0, k)
        return ans


if __name__ == '__main__':
    n1 = TreeNode(3)
    n2 = TreeNode(5)
    n3 = TreeNode(1)
    n4 = TreeNode(6)
    n5 = TreeNode(2)
    n6 = TreeNode(0)
    n7 = TreeNode(8)
    n8 = TreeNode(7)
    n9 = TreeNode(4)
    n1.left = n2
    n1.right = n3
    n2.left = n4
    n2.right = n5
    n3.left = n6
    n3.right = n7
    n5.left = n8
    n5.right = n9
    sol = Solution()
    print(sol.distanceK(n1, n2, 2))
